{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ec823036-79d6-4029-aeb8-636c6d663571",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "# Basic demonstration of Gemini SDK integration functionality with the Unity Catalog AI Toolkit SDK\n",
    "\n",
    "To get started with this, you will need a Google API Key. \n",
    "For testing purposes, you can generate a new account, and use your evaluation test key (no credit card required!).\n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `GOOGLE_API_KEY` after storing it in the `Databricks Secrets` API (accessible via dbutils or the databricks sdk workspace client)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c7a49c85-025f-4c28-b6ad-b4a0b302769d",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n",
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n",
      "\u001b[43mNote: you may need to restart the kernel using %restart_python or dbutils.library.restartPython() to use updated packages.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "%pip install -Uqq unitycatalog-ai unitycatalog-gemini\n",
    "%restart_python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "aa62833c-83ea-47a6-a50b-136ec4ab151e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Setting your API Key\n",
    "\n",
    "Don't forget to remove the key after you're done running cell 4!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "90d970f0-da57-44b0-816c-90fd9d0de55a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from databricks.sdk import WorkspaceClient\n",
    "\n",
    "workspace_client = WorkspaceClient(profile=\"optional to specify profile\")\n",
    "\n",
    "secret_scope = \"puneet_jain\"  # Change me!\n",
    "\n",
    "# Run this if you don't have the API key set to your secrets scope yet\n",
    "\n",
    "# if secret_scope not in [scope.name for scope in workspace_client.secrets.list_scopes()]:\n",
    "#     workspace_client.secrets.create_scope(secret_scope)\n",
    "\n",
    "# my_secret = \"<your API key, temporarily>\"\n",
    "\n",
    "# workspace_client.secrets.put_secret(scope=secret_scope, key=\"gemini_api_key\", string_value=my_secret)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fd54ca3c-b453-4ac6-968f-c47601ca91c1",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Fetch the key and set it to the environment variable key that the gemini SDK needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2f2ef55b-d3ba-4ddb-af72-ef110d8b9a2e",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "os.environ[\"GOOGLE_API_KEY\"] = base64.b64decode(\n",
    "    workspace_client.secrets.get_secret(scope=secret_scope, key=\"gemini_api_key\").value\n",
    ").decode()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "43b9d96b-1fc1-483c-93c2-ae9c33461790",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "assert \"GOOGLE_API_KEY\" in os.environ, (\n",
    "    \"Please set the GOOGLE_API_KEY environment variable to your Gemini API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "433bb01a-7d2b-44c5-ac32-0f0e8b71386a",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Import the UC client for Databricks UC. \n",
    "This will allow for function creation through either the `create_function` API (requires the defined `sql_body` statement) or the `create_python_function` (requires a type-hint-applied and docstring commented python callable). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "17f134d4-0fe2-4479-8717-3ce183165623",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n",
    "from unitycatalog.ai.gemini.toolkit import UCFunctionToolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6b8a211e-b576-4193-a085-6639b39c5f6f",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Set the UC Catalog and Schema \n",
    "\n",
    "You must set both of these that you will be using to store and execute your function(s). If these do not exist, ensure that you create them first. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3c8d77d3-321a-49b7-a13b-9f4ee73ba726",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "CATALOG = \"puneetjain_uc\"  # Change me!\n",
    "SCHEMA = \"autogen_ucai\"  # Change me if you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "15c65ce9-0464-4a4d-921c-21c149d244f2",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "client = DatabricksFunctionClient(profile=\"optional to specify profile\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "478f153f-7106-4b54-940a-0f61fe6e5d79",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "### Define a Callable\n",
    "The requirements for the callable:\n",
    "\n",
    "**typing**\n",
    "\n",
    "Types **must** be supplied for both the arguments and the return type. Function signatures that do not have these defined will raise a `ValueError`.\n",
    "\n",
    "The following types are not allowed:\n",
    "`Union`\n",
    "`Any`\n",
    "\n",
    "Additional caveats:\n",
    "Collections **must** supply typing of the interior components. For instance, ``typing.Dict`` is not allowed, but ``typing.Dict[str, str]`` will work correctly. \n",
    "\n",
    "**doc strings**\n",
    "\n",
    "The doc string **must** be in the Google Docstring format.\n",
    "Args and Returns comments are optional, but the function description **is required**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c85dde3b-626a-4f8b-9ff1-ac6e850469db",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def fetch_weather(location: str) -> str:\n",
    "    \"\"\"\n",
    "    Fetches the current weather in celsius for a given location.\n",
    "\n",
    "    Args:\n",
    "        location (str): The location to fetch the weather for.\n",
    "\n",
    "    Returns:\n",
    "        str: The current weather in celsius for the given location.\n",
    "    \"\"\"\n",
    "\n",
    "    return \"243.9 C\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f1d35198-5d37-4dda-9602-cab0648c3afd",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(browse_only=None, catalog_name='puneetjain_uc', comment='Fetches the current weather in celsius for a given location.', created_at=1734574994989, created_by='puneet.jain@databricks.com', data_type=<ColumnTypeName.STRING: 'STRING'>, external_language='Python', external_name=None, full_data_type='STRING', full_name='puneetjain_uc.autogen_ucai.fetch_weather', function_id='1501f129-50de-400e-b07f-898fe1754bf5', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='location', type_text='string', type_name=<ColumnTypeName.STRING: 'STRING'>, position=0, comment='The location to fetch the weather for.', parameter_default=None, parameter_mode=None, parameter_type=<FunctionParameterType.PARAM: 'PARAM'>, type_interval_type=None, type_json='{\"name\":\"location\",\"type\":\"string\",\"nullable\":true,\"metadata\":{\"comment\":\"The location to fetch the weather for.\"}}', type_precision=0, type_scale=0)]), is_deterministic=False, is_null_call=None, metastore_id='19a85dee-54bc-43a2-87ab-023d0ec16013', name='fetch_weather', owner='puneet.jain@databricks.com', parameter_style=<FunctionInfoParameterStyle.S: 'S'>, properties='{\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.taskWaitTimeInSeconds\":\"1000\",\"sqlConfig.spark.sql.ansi.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.decimal.dataType.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.statefulOperator.stateRebalancing.enabled\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.debugLogEnabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.remoteHttpClient.maxConnections\":\"2048\",\"sqlConfig.spark.sql.legacy.createHiveTableByDefault\":\"false\",\"sqlConfig.spark.sql.shuffleDependency.skipMigration.enabled\":\"true\",\"sqlConfig.spark.thriftserver.cloudStoreBasedRowSet.executor.cloudUploadThreshold\":\"6291456\",\"sqlConfig.spark.sql.streaming.stopTimeout\":\"15s\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.remoteHttpClient.timeoutInSeconds\":\"360\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdCurrentQpsIncreaseRatio\":\"0.0\",\"sqlConfig.spark.databricks.sql.functions.builtInHttpReq.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdTotalQpsIncreaseRatio\":\"0.0\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.useDedicatedHttpClient\":\"true\",\"sqlConfig.spark.sql.readSideCharPadding\":\"true\",\"sqlConfig.spark.sql.variable.substitute\":\"false\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.threadKeepAliveTimeInSeconds\":\"600\",\"sqlConfig.spark.databricks.sql.expression.aiFunctions.repartition\":\"0\",\"sqlConfig.spark.databricks.sql.functions.aiForecast.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.useDynamicTaskQueueExecutor\":\"false\",\"sqlConfig.spark.sql.sources.default\":\"delta\",\"sqlConfig.spark.sql.hive.convertCTAS\":\"true\",\"sqlConfig.spark.databricks.sql.functions.vectorSearch.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.safe.inference.enabled\":\"true\",\"sqlConfig.spark.sql.scripting.enabled\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batch.aiQuery.embedding.request.size\":\"4\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.batch.execution.size\":\"2048\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.clusterSizeBasedGlobalParallelism.scaleFactor\":\"32.0\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.embeddingsEndpointName\":\"databricks-gte-large-en\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.maxPoolSize\":\"2048\",\"sqlConfig.spark.sql.sources.commitProtocolClass\":\"com.databricks.sql.transaction.directory.DirectoryAtomicCommitProtocol\",\"sqlConfig.spark.sql.functions.remoteHttpClient.retryOn400TimeoutError\":\"true\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.adaptiveThreadPool.scaleUpThresholdSuccessRatio\":\"0.95\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.modelEndpointTypeParsing.enabled\":\"true\",\"sqlConfig.spark.sql.stableDerivedColumnAlias.enabled\":\"true\",\"sqlConfig.spark.sql.parquet.compression.codec\":\"snappy\",\"sqlConfig.spark.databricks.sql.functions.aiFunctions.model.parameters.enabled\":\"true\",\"sqlConfig.spark.sql.streaming.stateStore.providerClass\":\"com.databricks.sql.streaming.state.RocksDBStateStoreProvider\"}', return_params=None, routine_body=<FunctionInfoRoutineBody.EXTERNAL: 'EXTERNAL'>, routine_definition='\\n    return \"243.9 C\"\\n', routine_dependencies=None, schema_name='autogen_ucai', security_type=<FunctionInfoSecurityType.DEFINER: 'DEFINER'>, specific_name='fetch_weather', sql_data_access=<FunctionInfoSqlDataAccess.NO_SQL: 'NO_SQL'>, sql_path=None, updated_at=1734574994989, updated_by='puneet.jain@databricks.com')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(func=fetch_weather, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ef1f314c-aebd-4014-ac19-f5841fe00902",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Create a tool instance to use with Gemini\n",
    "\n",
    "toolkit = UCFunctionToolkit(function_names=[f\"{CATALOG}.{SCHEMA}.fetch_weather\"], client=client)\n",
    "\n",
    "tools = toolkit.generate_callable_tool_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fd08ff30-3b24-4abb-adb2-8841260507e7",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Submit the question\n",
    "\n",
    "In the question request, submit the defined tools from the `UCFunctionToolkit` instance functions that have been fetched. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "071e0891-23ae-4afc-bed8-e0d7ca43b18b",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "response:\n",
       "GenerateContentResponse(\n",
       "    done=True,\n",
       "    iterator=None,\n",
       "    result=protos.GenerateContentResponse({\n",
       "      \"candidates\": [\n",
       "        {\n",
       "          \"content\": {\n",
       "            \"parts\": [\n",
       "              {\n",
       "                \"text\": \"The weather in Nome, AK is 243.9 C, and the weather in Death Valley, CA is 243.9 C.\\n\"\n",
       "              }\n",
       "            ],\n",
       "            \"role\": \"model\"\n",
       "          },\n",
       "          \"finish_reason\": \"STOP\",\n",
       "          \"safety_ratings\": [\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_HARASSMENT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            }\n",
       "          ],\n",
       "          \"avg_logprobs\": -0.02099920402873646\n",
       "        }\n",
       "      ],\n",
       "      \"usage_metadata\": {\n",
       "        \"prompt_token_count\": 169,\n",
       "        \"candidates_token_count\": 33,\n",
       "        \"total_token_count\": 202\n",
       "      }\n",
       "    }),\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Interface with Gemini via their SDK\n",
    "from google.generativeai import GenerativeModel\n",
    "\n",
    "multi = \"What's the weather in Nome, AK and in Death Valley, CA?\"\n",
    "\n",
    "\n",
    "model = GenerativeModel(model_name=\"gemini-2.0-flash-exp\", tools=tools)\n",
    "\n",
    "chat = model.start_chat(enable_automatic_function_calling=True)\n",
    "\n",
    "response = chat.send_message(\"What's the weather in Nome, AK and in Death Valley, CA?\")\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e7f9c969-62ff-4372-bc99-a86c9293c7bb",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "source": [
    "## Show Details of Tool Calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "60dbf19e-c133-45c3-9c09-8605f8e61b40",
     "showTitle": false,
     "tableResultSettingsMap": {},
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user -> [{'text': \"What's the weather in Nome, AK and in Death Valley, CA?\"}]\n",
      "--------------------------------------------------------------------------------\n",
      "model -> [{'function_call': {'name': 'fetch_weather', 'args': {'location': 'Nome, AK'}}}, {'function_call': {'name': 'fetch_weather', 'args': {'location': 'Death Valley, CA'}}}]\n",
      "--------------------------------------------------------------------------------\n",
      "user -> [{'function_response': {'name': 'fetch_weather', 'response': {'result': '{\"format\": \"SCALAR\", \"value\": \"243.9 C\"}'}}}, {'function_response': {'name': 'fetch_weather', 'response': {'result': '{\"format\": \"SCALAR\", \"value\": \"243.9 C\"}'}}}]\n",
      "--------------------------------------------------------------------------------\n",
      "model -> [{'text': 'The weather in Nome, AK is 243.9 C, and the weather in Death Valley, CA is 243.9 C.\\n'}]\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "for content in chat.history:\n",
    "    print(content.role, \"->\", [type(part).to_dict(part) for part in content.parts])  # noqa: T201\n",
    "    print(\"-\" * 80)  # noqa: T201"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manually Execute Function calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "response:\n",
       "GenerateContentResponse(\n",
       "    done=True,\n",
       "    iterator=None,\n",
       "    result=protos.GenerateContentResponse({\n",
       "      \"candidates\": [\n",
       "        {\n",
       "          \"content\": {\n",
       "            \"parts\": [\n",
       "              {\n",
       "                \"text\": \"The weather in Nome, AK is 243.9 C and the weather in Death Valley, CA is 243.9 C.\\n\"\n",
       "              }\n",
       "            ],\n",
       "            \"role\": \"model\"\n",
       "          },\n",
       "          \"finish_reason\": \"STOP\",\n",
       "          \"safety_ratings\": [\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_HARASSMENT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            },\n",
       "            {\n",
       "              \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n",
       "              \"probability\": \"NEGLIGIBLE\"\n",
       "            }\n",
       "          ],\n",
       "          \"avg_logprobs\": -0.033611785620450974\n",
       "        }\n",
       "      ],\n",
       "      \"usage_metadata\": {\n",
       "        \"prompt_token_count\": 169,\n",
       "        \"candidates_token_count\": 32,\n",
       "        \"total_token_count\": 201\n",
       "      }\n",
       "    }),\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from google.generativeai.types import content_types\n",
    "\n",
    "from unitycatalog.ai.gemini.utils import generate_tool_call_messages, get_function_calls\n",
    "\n",
    "history = []\n",
    "question = \"What's the weather in Nome, AK and in Death Valley, CA?\"\n",
    "\n",
    "\n",
    "content = content_types.to_content(question)\n",
    "if not content.role:\n",
    "    content.role = \"user\"\n",
    "\n",
    "history.append(content)\n",
    "\n",
    "response = model.generate_content(history)\n",
    "while function_calls := get_function_calls(response):\n",
    "    history, function_calls = generate_tool_call_messages(\n",
    "        model=model, response=response, conversation_history=history\n",
    "    )\n",
    "\n",
    "    response = model.generate_content(history)\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": {
    "base_environment": "",
    "client": "1"
   },
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "Anthropic UC SDK tool usage",
   "widgets": {}
  },
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
